# 最常见的文档链，将文档直接塞入Prompt中，为LLM回答问题提供上下文资料，适合小文档场景
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.llm import LLMChain
from langchain_community.chat_models import ChatOpenAI
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.prompts import PromptTemplate
import __init__

loader = PyPDFLoader("../kecheng源码/loader.pdf")
print(loader.load())

prompt_template = """对以下文字做简洁的总结：
{text}
简洁的总结："""

prompt = PromptTemplate.from_template(prompt_template)
llm = ChatOpenAI(
    temperature=0,
    model="gpt-4-1106-preview",
)
llm_chain = LLMChain(llm=llm, prompt=prompt)

stuff_chain = StuffDocumentsChain(
    llm_chain=llm_chain,
    document_variable_name="text",
)
docs = loader.load()
print(stuff_chain.run(docs))